import numpy as np
import random
import torch

class EA:

    def __init__(
        self, 
        dimesion, 
        epsilon, 
        norm_type, 
        method, 
        evaluation, 
        config, 
        optimizer='GA', 
        strategy='random', 
        population_size=20, 
        perturbations=None, 
        isExtension=False,
    ):
        """
        给定约束，然后初始化策略，种群大小都要给出
        
        strategy: random: 随机在限制内生成种群, base: 基于基样本生成
        population_size: 配合strategy策略里的random策略使用
        perturbations: 如果strategy策略是base，则需要按照perturbations为基本去生成样本(分为仅以这些样本为种群去生成，还是再围绕这些样本去合成一些样本(可以配合population_size去使用))
        isExtension: 如果基于扰动生成人口，需要扩展一些人口的话，可以使用isExtension参数，初步是配合population_size参数生成满足种群个数的样本，满足\sum w_i = 1,目前使用随机生成
        epsilon: 规定范数的上限
        norm_type: L_inf还是L_2范数的攻击方式
        method: 是浮点数攻击还是整数攻击
        """
        self.dimesion = dimesion
        self.epsilon = epsilon
        self.norm_type = norm_type
        self.method = method
        self.strategy = strategy
        self.population_size = population_size
        self.perturbations = perturbations
        self.isExtension = isExtension
        self.evaluation = evaluation
        self.optimizer = optimizer
        self.config = config

        if self.config['dataset'] == 'CIFAR-10':
            self.mutation_rate = 1 / 3072
        elif self.config['dataset'] == 'MNIST' or self.config['dataset'] == 'FMNIST':
            self.mutation_rate = 1 / 784
        elif self.config['dataset'] == 'ImageNet':
            self.mutation_rate = 1 / 5000
        

        if method == 'float' and norm_type == 'inf':    ## 浮点数攻击， 无穷范数攻击
            
            if strategy == 'random':
                self.populations = epsilon * (-1 + 2 * np.random.rand(population_size, dimesion))

            elif strategy == 'base':
                if perturbations == None:
                    raise Exception('The init populations way is base, but the basesample is empty')

                if isExtension: ## 现在的策略是随机一次权重吗，和为1即可  TODO 可以改进这种策略
                    perturbations = torch.cat(perturbations, dim=0)  ## 添加一维，然后变成tensor张量                    
                    bs = perturbations.clone().reshape(perturbations.shape[0], dimesion).numpy()  ## 转成一维，然后转成numpy
                    size = bs.shape[0]
                    if size >= population_size:   ## 如果满足population_size的需求，就不添加了
                        self.populations = bs
                    else:
                        self.populations = np.zeros((population_size, dimesion))
                        for index in range(size):
                            self.populations[index] = bs[index]
                        for index in range(size, population_size):
                            weight = np.random.rand(size)
                            weight = weight / weight.sum()  ## 归一化 ,矩阵向量乘积，这里可以画图展示一下
                            self.populations[index] = np.dot(weight, bs)
                else:
                    perturbations = torch.cat(perturbations, dim=0)  ## 连接batch维度，变成tensor张量                    
                    self.populations = perturbations.clone().reshape(perturbations.shape[0], dimesion).numpy()  ## 转成一维，然后转成numpy

        elif method == 'float' and norm_type == 'l2':
            
            pass
        
        elif method == 'integer' and norm_type == 'inf':

            ## 整型编码工作， {-\epsilon, 0, \epsilon}
            if strategy == 'random':
                self.populations = epsilon * np.random.randint(low=-1, high=2, size=(population_size, dimesion))
            elif strategy == 'base' and self.config['dataset'] != 'ImageNet':
                if perturbations == None:
                    raise Exception('The init populations way is base, but the base sample is empty')

                if isExtension: ## 现在的策略是随机一次权重吗，和为1即可  TODO 可以改进这种策略
                    perturbations = torch.cat(perturbations, dim=0)  ## 添加一维，然后变成tensor张量                    
                    bs = perturbations.clone().reshape(perturbations.shape[0], dimesion).numpy()  ## 转成一维，然后转成numpy
                    size = bs.shape[0]
                    if size >= population_size:   ## 如果满足population_size的需求，就不添加了
                        self.populations = bs
                    else:
                        self.populations = np.zeros((population_size, dimesion))
                        for index in range(size):
                            self.populations[index] = bs[index]
                        for index in range(size, population_size):
                            weight = np.random.rand(size)
                            weight = weight / weight.sum()  ## 归一化 ,矩阵向量乘积，这里可以画图展示一下
                            self.populations[index] = epsilon * np.sign(np.dot(weight, bs))

            elif strategy == 'base' and self.config['dataset'] == 'ImageNet':
                print('Strategy for ImageNet')
                if perturbations == None:
                    raise Exception('The init populations way is base, but the basesample is empty')
                
                if isExtension: ## 现在的策略是随机一次权重吗，和为1即可  TODO 可以改进这种策略
                    perturbations = torch.cat(perturbations, dim=0)  ## 添加一维，然后变成tensor张量                    
                    bs = perturbations.clone().reshape(perturbations.shape[0], dimesion).numpy()  ## 转成一维，然后转成numpy
                    size = bs.shape[0]
                    if size >= population_size:   ## 如果满足population_size的需求，就不添加了
                        self.populations = bs
                    else:
                        self.populations = np.zeros((population_size, dimesion))
                        for index in range(size):
                            self.populations[index] = bs[index]
                        self.populations[size : population_size] = (epsilon * np.random.randint(low=-1, high=2, size=(population_size - size, dimesion)))
                else:
                    perturbations = torch.cat(perturbations, dim=0)  ## 连接batch维度，变成tensor张量                    
                    self.populations = perturbations.clone().reshape(perturbations.shape[0], dimesion).numpy()  ## 转成一维，然后转成numpy  
        
        elif method == 'integer' and norm_type == 'l2':
        
            pass
        
        else:
            raise ValueError

        self.population_size = self.populations.shape[0]

        self.fitness = self.evaluation(self.populations)
    
    def step(self, ):

        if self.optimizer == 'GA':
            if self.method == 'float':
                return self.GA_optimizer()
            else:
                return self.sign_GA_optimizer()
        else:
            if self.method == 'float':
                return self.DE_optimizer()
            else:
                return self.sign_DE_optimizer()

    ## 现在优化器是按照无穷范数的模式构造的，二范数的话可能需要改写一下， TODO
    ## 这四个生成方法像是优化器一样，实现内部实现迭代一次的优化步骤，返回种群和适应值
    def GA_optimizer(self, ):

        child_populations = np.zeros_like(self.populations)        
        
        ## selection
        parents = np.zeros((self.population_size * 2, self.dimesion))
        for i in range(2 * self.population_size):
            candidate_1 = random.randint(0, self.population_size - 1)
            candidate_2 = random.randint(0, self.population_size - 1)
            while candidate_1 == (i // 2) or candidate_2 == (i // 2) or candidate_1 == candidate_2:
                candidate_1 = random.randint(0, self.population_size - 1)
                candidate_2 = random.randint(0, self.population_size - 1)
            if self.fitness[candidate_1] > self.fitness[candidate_2]:
                parents[i] = self.populations[candidate_1]
            else:
                parents[i] = self.populations[candidate_2]
        ## cross
        for i in range(self.population_size):
            for j in range(self.dimesion):
                p = random.random()
                if p < 0.5:
                    child_populations[i, j] = parents[2 * i, j]
                else:
                    child_populations[i, j] = parents[2 * i + 1, j]
        ## mutation
        for i in range(self.population_size):
            for j in range(self.dimesion):
                p  = random.random()
                if p < self.mutation_rate:
                    child_populations[i, j] += self.config['alpha'] * self.epsilon * (-1 + 2 * random.random())
        
        ## clip
        child_populations = np.clip(child_populations, a_min=-self.epsilon, a_max=self.epsilon)
        ## evaluation
        child_fitness = self.evaluation(child_populations)

        for i in range(self.population_size):
            if child_fitness[i] > self.fitness[i]:
                self.populations[i] = child_populations[i]
                self.fitness[i] = child_fitness[i]
         
        
    def DE_optimizer(self, ):
        
        ## mutation
        Mpopulations = np.zeros_like(self.populations)
        for i in range(self.population_size):
            r1 = r2 = r3 = 0
            while r1 == i or r2 == i or r3 == i or r2 == r1 or r3 == r1 or r3 == r2:
                r1 = random.randint(0, self.population_size - 1)
                r2 = random.randint(0, self.population_size - 1)
                r3 = random.randint(0, self.population_size - 1)
        Mpopulations[i] = self.populations[r1] + self.config['F'] * (self.populations[r2] - self.populations[r3])
        Mpopulations[i] = np.clip(Mpopulations[i], a_min=-self.epsilon, a_max=self.epsilon)
        
        ## cross
        Cpopulations = np.zeros_like(self.populations)
        for i in range(self.population_size):
            for j in range(self.dimesion):
                rand_j = random.randint(0, self.dimesion - 1)
                rand_float = random.random()
                if rand_float <= self.config['CR'] or rand_j == j:
                    Cpopulations[i, j] = Mpopulations[i, j]
                else:
                    Cpopulations[i, j] = self.populations[i, j]

        ## selection
        Cfitness = self.evaluation(Cpopulations)
        for i in range(self.population_size):
            if Cfitness[i] > self.fitness[i]:
                self.populations[i] = Cpopulations[i]
                self.fitness[i] = Cfitness[i]

    def sign_sampling(self, value):

        integer_value = np.sign(value)
        instead_integer_value = np.random.randint(low=-1, high=2)
        while integer_value == instead_integer_value:
            instead_integer_value = np.random.randint(low=-1, high=2)
        return instead_integer_value

    def sign_GA_optimizer(self, ):

        child_populations = np.zeros_like(self.populations)

        ## selection
        parents = np.zeros((self.population_size * 2, self.dimesion))
        for i in range(2 * self.population_size):
            candidate_1 = random.randint(0, self.population_size - 1)
            candidate_2 = random.randint(0, self.population_size - 1)
            while candidate_1 == (i // 2) or candidate_2 == (i // 2) or candidate_1 == candidate_2:
                candidate_1 = random.randint(0, self.population_size - 1)
                candidate_2 = random.randint(0, self.population_size - 1)
            if self.fitness[candidate_1] > self.fitness[candidate_2]:
                parents[i] = self.populations[candidate_1]
            else:
                parents[i] = self.populations[candidate_2]

        ## crossover
        for i in range(self.population_size):
            for j in range(self.dimesion):
                p = random.random()
                if p < 0.5:
                    child_populations[i, j] = parents[2 * i, j]
                else:
                    child_populations[i, j] = parents[2 * i + 1, j]

        ## mutation
        for i in range(self.population_size):
            for j in range(self.dimesion):
                p = random.random()
                if p < self.mutation_rate:
                    child_populations[i, j] = self.epsilon * self.sign_sampling(child_populations[i, j])
        ## clip 
        child_populations = np.clip(child_populations, a_min=-self.epsilon, a_max=self.epsilon)

        ## evaluation
        child_fitness = self.evaluation(child_populations)

        for i in range(self.population_size):
            if child_fitness[i] > self.fitness[i]:
                self.populations[i] = child_populations[i]
                self.fitness[i] = child_fitness[i]

    def sign_DE_optimizer(self, ):

        pass